import tensorflow as tf
import load_data
import data_set
import base_model
import ufans_model
import config as cfg
import os
import numpy as np


def pad_x(x_a):
    num = 2
    l = x_a.shape[1]
    while l > num:
        num *= 2
    pad_len = num - l
    x_pad = np.zeros((x_a.shape[0], pad_len, x_a.shape[2]))
    return np.concatenate((x_a, x_pad), axis=1)


class ModelRun(object):

    def __init__(self, model):
        self.model = model
        self.saver = tf.train.Saver()
        self.sess = None
        self.x_d_min = None
        self.x_d_max = None

    def train_duration(self, d_set):
        train_iter = d_set.train_gen()
        sess = self.sess
        model = self.model
        sess.run(train_iter.initializer)
        next_item = train_iter.get_next()
        for i in range(1, cfg.d_train_steps + 1):
            model.global_step += 1
            model.d_step += 1
            x_d, y_d, ul_d, x_a, y_a, ul_a = sess.run(next_item)
            x_a = pad_x(x_a)
            y_a = pad_x(y_a)
            if model.global_step % cfg.log_steps == 0:
                _, d_loss, a_loss = sess.run((model.d_train_op,
                                              model.d_loss,
                                              model.a_loss,
                                              ),
                                             feed_dict={
                                                 model.x_d_ph: x_d,
                                                 model.y_d_ph: y_d,
                                                 model.lu_d_ph: ul_d,
                                                 model.x_a_ph: x_a,
                                                 model.y_a_ph: y_a,
                                                 model.lu_a_ph: ul_a
                                             })
                print('step %d , d_loss %f, a_loss %f' % (
                    model.global_step, d_loss, a_loss))
            else:
                sess.run(model.d_train_op, feed_dict={
                    model.x_d_ph: x_d,
                    model.y_d_ph: y_d,
                    model.lu_d_ph: ul_d,
                    model.x_a_ph: x_a,
                    model.y_a_ph: y_a,
                    model.lu_a_ph: ul_a
                })

        if model.global_step % cfg.save_steps == 0:
            self.saver.save(self.sess, cfg.log_dir)

    def train_acoustic(self, d_set):
        train_iter = d_set.train_gen()
        sess = self.sess
        model = self.model
        sess.run(train_iter.initializer)
        next_item = train_iter.get_next()
        for i in range(1, cfg.d_train_steps + 1):
            model.global_step += 1
            model.a_step += 1
            x_d, y_d, ul_d, x_a, y_a, ul_a = sess.run(next_item)
            x_a = pad_x(x_a)
            y_a = pad_x(y_a)
            if model.global_step % cfg.log_steps == 0:
                _, d_loss, a_loss = sess.run((model.a_train_op,
                                              model.d_loss,
                                              model.a_loss,
                                              ),
                                             feed_dict={
                                                 model.x_d_ph: x_d,
                                                 model.y_d_ph: y_d,
                                                 model.lu_d_ph: ul_d,
                                                 model.x_a_ph: x_a,
                                                 model.y_a_ph: y_a,
                                                 model.lu_a_ph: ul_a
                                             })
                print('step %d , d_loss %f, a_loss %f' % (
                    model.global_step, d_loss, a_loss))
            else:

                sess.run(model.a_train_op, feed_dict={
                    model.x_d_ph: x_d,
                    model.y_d_ph: y_d,
                    model.lu_d_ph: ul_d,
                    model.x_a_ph: x_a,
                    model.y_a_ph: y_a,
                    model.lu_a_ph: ul_a
                })

        if model.global_step % cfg.save_steps == 0:
            self.saver.save(self.sess, cfg.log_dir)

    def run(self):
        ld = load_data.LoadedData()
        ld.load()
        ld.norm()
        d_set = data_set.DataSet(ld, cfg.batch_size)
        with tf.Session() as sess:
            self.sess = sess
            sess.run(tf.global_variables_initializer())
            # self.train_duration(d_set)
            self.train_acoustic(d_set)


if __name__ == '__main__':
    os.makedirs(cfg.log_dir, exist_ok=True)
    os.environ['CUDA_VISIBLE_DEVICES'] = cfg.gpu
    model_ = ufans_model.UFANSModel()
    # model_ = base_model.BaseTTSModel()
    model_.build_graph()
    model_run = ModelRun(model_)
    model_run.run()
